import torch
import numpy as np
import matplotlib.pyplot as plt


class FFTBlock:
    def __init__(self, topk):
        self.topk = topk

    '''
    根据输入的周期成分，利用傅里叶分解找出其振幅和频率
    利用振幅大小挑选前k个最合适的周期
    输入是(batch, time_seq, features)
    返回振幅是(batch, top_k)，周期是(top_k)
    '''
    def decomp(self, season:torch.Tensor):
        # B, T, N分别表示样本数，时间序列的长度和特征数
        B, T, N = season.shape

        # 傅里叶分解得到频域序列，并计算每个点对应的频率
        fft_result = torch.fft.rfft(season, dim=1)
        freq = torch.fft.rfftfreq(T, d=1.0).to(season.device)
        
        # 计算每个频率对应的振幅，对每个样本，计算所有特征的平均振幅，为了比较振幅最大的周期，取所有样本振幅的和ampsum用来比较
        amplitudes = torch.mean(torch.abs(fft_result), dim=-1, keepdim=False)
        ampsum = torch.sum(amplitudes, dim=0, keepdim=False)
        
        # 忽略直流分量
        ampsum[0] = 0
        
        # 选出振幅最大的k个振幅的下标
        _, topk_indice = torch.topk(ampsum, self.topk)
        
        # 返回振幅（用作后续softmax），周期（和挑选出的振幅同坐标的频率，再取倒数就是周期）
        return amplitudes[:, topk_indice], (1 / freq[topk_indice]).clone().detach().to(torch.int)
    

if __name__ == "__main__":
    # 随机生成一个时间序列
    torch.manual_seed(0)
    B = 112
    T = 32
    C = 3
    X = torch.randn(B, T, C)
    fftblock = FFTBlock(topk=5)
    amp, frep = fftblock.decomp(X)
    print(amp)
    print(frep)



